package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/fatih/color"
	"golang.org/x/sync/errgroup"
	"golang.org/x/text/encoding/simplifiedchinese"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync/atomic"
	"time"
)

const Version = "1.3.0"

// Protoc 指定版本，避免生成文件差异及实现多版本生成
const Protoc = "protoc-3.19.1"

var (
	// 忽略的路径，相对于proto目录，包含匹配，路径分隔符统一/
	ignoreDir = []string{"google/"}
	// 忽略替换的路径，包含匹配，路径分隔符统一/
	ignoreReplaceOmitempty []string
	// 自动发现子模块路径，多个用逗号分隔
	autoDiscoverySubmoduleProto = "proto,common"
	// proto子模块根路径
	basePath string

	// 老版本proto
	isOld bool
	// 获取proto方式
	findMethod string
	// 通过差异对比时，当前分支前N个接口
	beforeNode int
)

func init() {
	basePath, _ = os.Getwd()
	// 如果当前目录下自动发现目录存在
	// 且是存在.git文件，则设置子模块为proto目录
	for _, s := range strings.Split(autoDiscoverySubmoduleProto, ",") {
		if _, err := os.Stat(s + string(os.PathSeparator) + ".git"); err != nil {
			continue
		} else {
			basePath += string(os.PathSeparator) + s
			break
		}
	}
}

func main() {
	var isVersion bool

	flag.BoolVar(&isVersion, "v", false, "查看版本")

	flag.StringVar(&findMethod, "m", "d", "获取proto方式，d差异文件（通过git diff HEAD~n实现）、a所有文件、s变更未提交的文件（通过git status实现）")
	flag.IntVar(&beforeNode, "n", 5, "通过差异（git diff HEAD~n）对比获取proto时，当前分支前N个节点")
	flag.BoolVar(&isOld, "o", true, "是否是旧版本")

	flag.Parse()

	if isVersion {
		fmt.Println(Version)
		return
	}

	var ps []string
	var err error

	// 记录开始时间
	startTime := time.Now()

	switch findMethod {
	case "a":
		ps, err = getProtoByDir(basePath)
	case "s":
		ps, err = getProtoByGit("git", "status", "--porcelain", "*.proto")
	case "d":
		ps, err = getProtoByGit("git", "-P", "diff", fmt.Sprintf("HEAD~%d", beforeNode), "--name-status", "*.proto")
	default:
		err = errors.New("参数错误，通过 -h 参数查看使用说明")
	}

	if err != nil {
		color.Red("》" + err.Error())
		return
	}

	color.Green("》找到%d个文件\n", len(ps))

	if err = buildMultiProto(ps); err != nil {
		color.Red("》" + err.Error())
		return
	}

	color.Green("\n》%d 个文件处理成功，共耗时 %d 毫秒\n\n", len(ps), time.Now().Sub(startTime)/time.Millisecond)
	return
}

// 通过git命令获取文件
func getProtoByGit(name string, args ...string) (ps []string, err error) {
	// 先检测proto目录是否启用git
	if _, err = os.Stat(basePath + string(os.PathSeparator) + ".git"); err != nil {
		if os.IsNotExist(err) {
			return nil, errors.New("执行方式依赖git命令，当前目录未启用git，可使用全量生成参数 -m a，详情查看帮助 -h")
		} else {
			return
		}
	}

	args = append([]string{"-C", basePath}, args...)
	cmd := exec.Command(name, args...)

	var out string
	outBytes, err := cmd.CombinedOutput()
	if len(outBytes) > 0 {
		out = string(outBytes)
	}

	// 命令记录日志
	cmdString := name + " " + strings.Join(args, " ")

	if err != nil {
		return nil, errors.New(fmt.Sprintf(`执行git命令 %s 出错 %s
%s`, cmdString, err.Error(), out))
	} else if len(out) > 0 {
		color.Blue("》执行git命令 %s 成功，结果：\n%s\n", cmdString, out)
	} else {
		color.Blue("》执行git命令 %s 成功，无变更文件\n", cmdString)
		return
	}

	// 按换行分隔
	changes := strings.FieldsFunc(out, func(c rune) bool {
		return c == '\n' || c == '\r'
	})

	for _, change := range changes {
		ca := strings.FieldsFunc(strings.TrimSpace(change), func(c rune) bool {
			return c == ' ' || c == '\t'
		})
		// 删除的不处理
		if len(ca) < 2 || ca[0] == "D" || pathArrayContains(ca[1], ignoreDir) {
			continue
		}
		// 统一分路径分隔符
		// R092    mm/member.proto sh/member.proto
		relativePath := strings.ReplaceAll(ca[len(ca)-1], "/", string(os.PathSeparator))
		// 绝对路径
		ps = append(ps, basePath+string(os.PathSeparator)+relativePath)
	}

	return
}

// 获取指定目录下所有的proto
func getProtoByDir(dirname string) (ps []string, err error) {
	file, err := os.ReadDir(dirname)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("读取目录 %s 出错：%s", dirname, err.Error()))
	}
	for _, v := range file {
		if v.IsDir() {
			if !pathArrayContains(v.Name(), ignoreDir) {
				if subPs, err := getProtoByDir(dirname + string(os.PathSeparator) + v.Name()); err != nil {
					return nil, err
				} else {
					ps = append(ps, subPs...)
				}
			}
		} else {
			path := dirname + string(os.PathSeparator) + v.Name()
			if strings.HasSuffix(v.Name(), ".proto") && !pathArrayContains(path, ignoreDir) {
				ps = append(ps, path)
			}
		}
	}
	return
}

// 批量构建proto
func buildMultiProto(ps []string) error {
	if len(ps) == 0 {
		return nil
	}

	// 并发控制
	ch := make(chan struct{}, runtime.GOMAXPROCS(0))
	// 成功计数
	count := int32(0)

	var eg errgroup.Group

	for _, p := range ps {
		ch <- struct{}{}
		p := p
		eg.Go(func() error {
			defer func() {
				<-ch
			}()

			rf := strings.ReplaceAll(p, basePath+string(os.PathSeparator), "")

			err := buildProto(p)
			current := atomic.AddInt32(&count, 1)
			// 如果出错则退出
			if err != nil {
				return errors.New(fmt.Sprintf("【%d/%d】%s 处理失败，执行中断\n》%s\n", current, len(ps), rf, err.Error()))
			} else {
				fmt.Printf("【%d/%d】%s 处理成功\n", current, len(ps), rf)
			}

			return nil
		})
	}

	return eg.Wait()
}

// 构建proto
func buildProto(path string) error {
	relativePath := strings.ReplaceAll(path, basePath+string(os.PathSeparator), "")

	var args []string
	if isOld {
		args = []string{"--go-1.3.2_out=plugins=grpc:.", relativePath}
	} else {
		if strings.HasSuffix(basePath, "oms-proto") {
			args = []string{"--go-1.27.1_out=paths=source_relative:.", "--go-grpc-1.1.0_out=paths=source_relative,require_unimplemented_servers=false:.", relativePath}
		} else {
			dir := filepath.Dir(relativePath)
			args = []string{"--go-1.27.1_out=" + dir, "--go-grpc-1.1.0_out=require_unimplemented_servers=false:" + dir, relativePath}
		}
	}

	cmd := exec.Command(Protoc, args...)
	cmd.Dir = basePath

	var out string
	outBytes, err := cmd.CombinedOutput()
	if len(outBytes) > 0 {
		// windows平台中文会乱码，转一下
		if runtime.GOOS == "windows" {
			outBytes, _ = simplifiedchinese.GB18030.NewDecoder().Bytes(outBytes)
		}
		out = string(outBytes)
	}

	if err != nil {
		return errors.New(fmt.Sprintf(`执行 %s 命令出错 %s，
%s`, Protoc+" "+strings.Join(args, " "), err.Error(), out))
	} else if len(out) > 0 {
		fmt.Println("	" + out)
	}

	if err := formatPb(strings.ReplaceAll(path, ".proto", ".pb.go"), true); err != nil {
		return errors.New("formatPb出错：" + err.Error())
	}
	if !isOld {
		if err := formatPb(strings.ReplaceAll(path, ".proto", "_grpc.pb.go"), false); err != nil {
			return errors.New("formatPb出错：" + err.Error())
		}
	}
	return nil
}

// 格式化pb文件
// 1、替换默认生成的Omitempty
// 2、windows平台 lf替换为crlf
func formatPb(pb string, replaceOmit bool) error {
	if replaceOmit {
		replaceOmit = !pathArrayContains(pb, ignoreReplaceOmitempty)
	}
	// 如果是忽略替换omitempty同时是window，则表示不需要处理
	if !replaceOmit && runtime.GOOS != "windows" {
		return nil
	}

	// 文件不存在的忽略
	if f, err := os.Stat(pb); err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	} else {
		if data, err := os.ReadFile(pb); err != nil {
			return err
		} else {
			if replaceOmit {
				data = bytes.ReplaceAll(data, []byte(",omitempty"), []byte(""))
			}
			if runtime.GOOS == "windows" {
				data = bytes.ReplaceAll(data, []byte{'\n'}, []byte{'\r', '\n'})
			}
			if err = os.WriteFile(pb, data, f.Mode()); err != nil {
				return err
			}
		}
	}

	return nil
}

// 数组中值是否包含匹配字符串
func pathArrayContains(s string, arr []string) bool {
	// 统一分隔符
	s = strings.ReplaceAll(s, string(os.PathSeparator), "/")
	for _, a := range arr {
		if strings.Contains(s, a) {
			return true
		}
	}
	return false
}
